import websocket
import redis
import json
import requests
import numpy as np
from datetime import datetime

pool = redis.ConnectionPool(host='localhost',
                            port=6379,
                            decode_responses=True,
                            db=0)
r = redis.Redis(connection_pool=pool)


class ws_orderbook:

    def __init__(self, symbol, interval, n):
        self.bids = {}
        self.asks = {}
        self.symbol = symbol
        self.interval = interval
        self.n = n

    def get_rest_snapshot(self):
        url = f'https://api.binance.com/api/v3/depth?symbol={self.symbol.upper()}&limit=1000'
        res = requests.get(url)
        data = json.loads(res.text.encode())
        snapshot_updateId = data['lastUpdateId']
        return snapshot_updateId

    def on_message(self, ws, message):
        result = json.loads(message)
        ask = np.array(result['a']).astype(float)
        bid = np.array(result['b']).astype(float)
        E = result['E']
        ts = E / 1000
        t = datetime.fromtimestamp(ts)
        str_t = datetime.strftime(t, '%Y-%m-%d %H:%M:%S')

        if len(bid) > 0:
            bid = dict(zip(bid[:, 0], bid[:, 1]))
            self.bids.update(bid)

            bids = {k: v for k, v in self.bids.items() if v != 0}
            sorted_bid_price = sorted(self.bids.keys(), reverse=True)[:self.n]  # bids按价格从高到低排序取前n个
            bid_vols = [self.bids.get(x) for x in sorted_bid_price]

        if len(ask) > 0:
            ask = dict(zip(ask[:, 0], ask[:, 1]))
            self.asks.update(ask)

            asks = {k: v for k, v in self.asks.items() if v != 0}

            sorted_ask_price = sorted(self.asks.keys())[:self.n]  # asks按价格从低到高排序取前n个（后续再取返序::-1）使得从高到低排序
            ask_vols = [asks.get(x) for x in sorted_ask_price]
            sorted_ask_price = sorted_ask_price[::-1]
            ask_vols = ask_vols[::-1]

        # asks 是降序的, bids 是升序的
        sorted_ask_price = np.array(sorted_ask_price)
        sorted_bid_price = np.array(sorted_bid_price)
        ask_vols = np.array(ask_vols)
        bid_vols = np.array(bid_vols)

        depth_data = {'t': str_t,
                      'ask_price': list(sorted_ask_price), 'ask_vols': list(ask_vols),
                      'bid_price': list(sorted_bid_price), 'bid_vols': list(bid_vols)}
        print(depth_data)
        depth_data = json.dumps(depth_data)
        r.set(f'{self.symbol}_depth', depth_data)

    def on_error(self, ws, error):
        print(error)

    def on_close(self, close_msg):
        print("### closed ###" + close_msg)

    def streamKline(self):
        websocket.enableTrace(False)
        socket = f'wss://stream.binance.com:9443/ws/{self.symbol}@depth@{self.interval}'
        ws = websocket.WebSocketApp(socket,
                                    on_message=self.on_message,
                                    on_error=self.on_error,
                                    on_close=self.on_close)

        ws.run_forever()


ob = ws_orderbook('solusdt', '1000ms', 500)
ob.streamKline()
